//
//  Sample.cpp
//  RayTracer
//
//  Created by 俞云康 on 2/17/16.
//  Copyright © 2016 yushroom. All rights reserved.
//

#include "Sample.hpp"

Sampler::~Sampler() {}

Sampler::Sampler(int samplesPerPixel) : m_samplersPerPixel(samplesPerPixel) {
    
}

CameraSample Sampler::getCameraSample(const Point2i &pRaster) {
    CameraSample cs;
    cs.pFilm = Point2f(float(pRaster.x), float(pRaster.y)) + get2D();
    cs.time = get1D();
    cs.pLens = get2D();
    return cs;
}

bool Sampler::setSampleNumber(int sampleNumber) {
    m_array1DOffest = m_array2DOffset = 0;
    m_currentPixelSampleIndex = sampleNumber;
    return m_currentPixelSampleIndex < m_samplersPerPixel;
}

void Sampler::startPixel(const Point2i &p) {
    m_currentPixel = p;
    m_currentPixelSampleIndex = 0;
    m_array1DOffest = m_array2DOffset = 0;
}

bool Sampler::startNextSample() {
    m_array1DOffest = m_array2DOffset = 0;
    return ++m_currentPixelSampleIndex < m_samplersPerPixel;
}

void Sampler::request1DArray(int n) {
    Assert(n == roundCount(n));
    m_samples1DArraySizes.push_back(n);
    m_sampleArray1D.push_back(vector<float>(n*m_samplersPerPixel));
}

void Sampler::request2DArray(int n) {
    Assert(n == roundCount(n));
    m_samples2DArraySizes.push_back(n);
    m_sampleArray2D.push_back(vector<Point2f>(n*m_samplersPerPixel));
}

const float * Sampler::get1DArray(int n) {
    if (m_array1DOffest == m_sampleArray1D.size())
        return nullptr;
    Assert(n == m_samples1DArraySizes[m_array1DOffest]);
    Assert(m_currentPixelSampleIndex < m_samplersPerPixel);
    return &m_sampleArray1D[m_array1DOffest++][m_currentPixelSampleIndex * n];
}

const Point2f * Sampler::get2DArray(int n) {
    if (m_array2DOffset == m_sampleArray2D.size())
        return nullptr;
    Assert(m_currentPixelSampleIndex < m_samplersPerPixel);
    return &m_sampleArray2D[m_array2DOffset++][m_currentPixelSampleIndex * n];
}

//void GlobalSampler::startPixel(const Point2i &p) {
//    Sampler::startPixel(p);
//    m_dimension = 0;
//    m_intervalSampleIndex = getIndexForSample(0);
//    m_arrayEndDim = arrayStartDim + m_sampleArray1D.size() + 2 * m_sampleArray2D.size();
//    
//    for (size_t i = 0; i < m_samples1DArraySizes.size(); ++i) {
//        int nSamples = m_samples1DArraySizes[i] * m_samplersPerPixel;
//        for (int j = 0; j < nSamples; ++j) {
//            int index = getIndexForSample(j);
//            m_sampleArray1D[i][j] = sampleDimension(index, arrayStartDim + i);
//        }
//    }
//    
//    int dim = arrayStartDim + m_samples1DArraySizes.size();
//    for (size_t i = 0; i < m_samples2DArraySizes.size(); ++i) {
//        int nSamples = m_samples2DArraySizes[i] * m_samplersPerPixel;
//        for (int j = 0; j < nSamples; ++j) {
//            int idx = getIndexForSample(j);
//            m_sampleArray2D[i][j].x = sampleDimension(idx, dim);
//            m_sampleArray2D[i][j].y = sampleDimension(idx, dim+1);
//        }
//        dim += 2;
//    }
//    Assert(dim == m_arrayEndDim);
//}
//
//bool GlobalSampler::startNextSample() {
//    m_dimension = 0;
//    m_intervalSampleIndex = getIndexForSample(m_currentPixelSampleIndex + 1);
//    return Sampler::startNextSample();
//}
//
//bool GlobalSampler::setSampleNumber(int sampleNumber) {
//    m_dimension = 0;
//    m_intervalSampleIndex = getIndexForSample(sampleNumber);
//    return Sampler::setSampleNumber(sampleNumber);
//}
//
//float GlobalSampler::get1D() {
//    if (m_dimension >= arrayStartDim && m_dimension < m_arrayEndDim)
//        m_dimension = m_arrayEndDim;
//    return sampleDimension(m_intervalSampleIndex, m_dimension++);
//}
//
//Point2f GlobalSampler::get2D() {
//    if (m_dimension + 1 >= arrayStartDim && m_dimension < m_arrayEndDim)
//        m_dimension = m_arrayEndDim;
//    Point2f p(sampleDimension(m_intervalSampleIndex, m_dimension),
//              sampleDimension(m_intervalSampleIndex, m_dimension + 1));
//    m_dimension += 2;
//    return p;
//}


RandomSampler::RandomSampler(int samplesPerPixel, int seed) 
	: Sampler(samplesPerPixel), rng(seed) 
{
    m_currentPixelSampleIndex = 0;
}

float RandomSampler::get1D() {
    Assert(m_currentPixelSampleIndex < m_samplersPerPixel);
    return rng.UniformFloat();
}

Point2f RandomSampler::get2D() {
    Assert(m_currentPixelSampleIndex < m_samplersPerPixel);
    return {rng.UniformFloat(), rng.UniformFloat()};
}

unique_ptr<Sampler> RandomSampler::clone(int seed) {
    RandomSampler *rs = new RandomSampler(*this);
    rs->rng.SetSequence(seed);
    return unique_ptr<Sampler>(rs);
}

void RandomSampler::startPixel(const Point2i &p) {
    for (size_t i = 0; i < m_sampleArray1D.size(); ++i)
        for (size_t j = 0; j < m_sampleArray1D[i].size(); ++j)
            m_sampleArray1D[i][j] = rng.UniformFloat();
    
    for (size_t i = 0; i < m_sampleArray2D.size(); ++i)
        for (size_t j = 0; j < m_sampleArray2D[i].size(); ++j)
            m_sampleArray2D[i][j] = {rng.UniformFloat(), rng.UniformFloat()};
    Sampler::startPixel(p);
}


// see PBRT 667
void ConcentricSampleDisk(float u1, float u2, float*dx, float*dy) {
    float r, theta;
    // Map uniform random numbers to $[-1,1]^2$
    float sx = 2 * u1 - 1;
    float sy = 2 * u2 - 1;
    
    // Map square to $(r,\theta)$
    
    // Handle degeneracy at the origin
    if (sx == 0.0 && sy == 0.0) {
        *dx = 0.0;
        *dy = 0.0;
        return;
    }
    if (sx >= -sy) {
        if (sx > sy) {
            // Handle first region of disk
            r = sx;
            if (sy > 0.0) theta = sy/r;
            else          theta = 8.0f + sy/r;
        }
        else {
            // Handle second region of disk
            r = sy;
            theta = 2.0f - sx/r;
        }
    }
    else {
        if (sx <= sy) {
            // Handle third region of disk
            r = -sx;
            theta = 4.0f - sy/r;
        }
        else {
            // Handle fourth region of disk
            r = -sy;
            theta = 6.0f + sx/r;
        }
    }
    theta *= M_PI / 4.f;
    *dx = r * cosf(theta);
    *dy = r * sinf(theta);
}